#load required packages
library(doParallel)
library(vegan)

#set up more lightweight transmission function, based on the transmission function in HERAChp.KandlerCrema
transmission<-function(N = 500, timesteps = 501, mu = 0.01, warmUp = 300, top = NA, bias = 0, raw = FALSE){
  agents <- 1:N
  traitCounter <- N + 1
  obs.div <- c()
  for(t in 1:timesteps){ 	
    samplePool = agents
    if(length(unique(samplePool)) > 1){
      sampleTraits = as.numeric(names(table(samplePool)))
      sampleProp = table(samplePool)/sum(table(samplePool))
      sampleTraitsProb = sampleProp^(1 - bias)/sum(sampleProp^(1 - bias))
      agents = sample(sampleTraits, size = N, replace = TRUE, prob = sampleTraitsProb)
    }
    index = which(runif(N) <= mu)
    if(length(index) > 0){
      newTraits <- traitCounter:c(traitCounter + length(index) - 1)
      agents[index] = newTraits
      traitCounter = max(newTraits) + 1
    }
    obs.div[t] <- diversity(table(agents), index = "simpson")
  }
  return(list(obs.div = obs.div))
}

#set up parallel computing
cl <-  makeForkCluster(8)
registerDoParallel(cl)

#run simulations
nsim <- 1000
timesteps <- 500
eq_calculation <- foreach(i = 1:nsim, .combine = c) %dopar% {
  transmission(N = 729, timesteps = timesteps, mu = 0.03672527, warmUp = 0, top = NA, bias = 0, raw = FALSE)$obs.div
}

#stop parallel computing
stopCluster(cl)

#reformat results for plotting
eq_calculation_raw <- eq_calculation
eq_calculation <- matrix(eq_calculation, nrow = timesteps, ncol = nsim)
